import copy
import math

import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from torch.nn import functional as F


# embedding = nn.Embedding(10, 3)
# input1 = torch.LongTensor([[1, 2, 4, 5], [4, 3, 2, 9]])
# print(embedding(input1))
#
# embedding = nn.Embedding(10, 3, padding_idx=0)
# input1 = torch.LongTensor([[1, 0, 4, 0], [4, 3, 2, 9]])
# print(embedding(input1))


class Embeddings(nn.Module):
    def __init__(self, vocab, d_model):
        """

        :param d_model: 词嵌入的维度
        :param vocab: 词表的大小
        """
        super(Embeddings, self).__init__()
        # 定义Embedding层
        self.lut = nn.Embedding(vocab, d_model)
        # 将参数传入类中
        self.d_model = d_model

    def forward(self, x):
        # x: 代表输入进模型的文本通过词汇映射后的数字张量
        return self.lut(x) * math.sqrt(self.d_model)



class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=5000):
        """

        :param d_model: 词嵌入的维度
        :param dropout: Dropout比率
        :param max_len: 每句的最大长度
        """
        super(PositionalEncoding, self).__init__()
        # 初始化dropout层
        self.dropout = nn.Dropout(dropout)
        # 初始化一个位置编码矩阵，大小是max_len * d_model
        pe = torch.zeros(max_len, d_model)
        # 初始化一个绝对位置矩阵, max_len * 1
        position = torch.arange(0, max_len).unsqueeze(1)
        # 定义一个变换矩阵div_term,跳跃式的初始化
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
        # 将变换矩阵进行奇数和矩阵位置分别赋值
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        # 将二维张量扩展为三维张量
        pe = pe.unsqueeze(0)
        # 将位置编码矩阵注册成模型的buffer， 不是模型中的超参数，不跟随优化器同步更新
        # 注册成buffer后，我们就可以在模型保存后重新加载，将这个位置编码器和模型参数一共加载进来
        self.register_buffer('pe', pe)

    def forward(self, x):
        """

        :param x: 文本序列的词嵌入表示
        :return:
        """
        # pe的编码太长了，将max_len对应的维度缩小为x的句子的长度
        pe_output = self.pe[:, :x.size(1)]
        pe_output.requires_grad_(False)
        # pe_output = pe_output.unsqueeze(0)
        x += pe_output
        return self.dropout(x)


def draw_plt():
    plt.figure(figsize=(15, 5))
    pe = PositionalEncoding(20, 0)
    y = pe(torch.zeros(1, 100, 20))
    plt.plot(np.arange(100), y[0, :, 4:8].data.numpy())
    plt.legend(['dim %d' % p for p in [4, 5, 6, 7]])
    plt.show()
    plt.savefig('positional_encoding.png')


def subsequent_mask(size):
    """
    构建掩码张量的函数，(掩码张量一定是个方阵？)
    :param size: 方阵的维度
    :return:
    """
    attn_shape = (1, size, size)
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')

    # 使得这个三角矩阵反转
    return torch.from_numpy(1 - subsequent_mask)


def attention(query, key, value, mask=None, dropout=None):
    """
    实现注意力机制，query, key, value: 代表注意力的三个输入张量
    :param query:
    :param key:
    :param value:
    :param mask:
    :param dropout:
    :return:
    """
    # query的最后一个维度是词嵌入的维度
    d_k = query.size(-1)
    # 按照注意力计算公式，将query和key的转置进行矩阵乘法，然后除以缩放系数
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    # 判断是否使用掩码张量
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    # p_attn 是经过 softmax 归一化后的注意力权重（attention weights）。
    # 它表示每个位置对其他位置的关注程度。
    p_attn = F.softmax(scores, dim=-1)

    # 判断是否使用dropout
    if dropout is not None:
        p_attn = dropout(p_attn)

    # 最后一步完成p_attn和value张量的乘法, 并返回query注意力表示
    return torch.matmul(p_attn, value), p_attn


def clones(module: nn.Module, N):
    """
    用于生成相同网络层的克隆函数
    :param module:要克隆的目标网络层
    :param N:需要克隆的数量
    :return:
    """
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


class MultiHeadedAttention(nn.Module):
    """

    """
    def __init__(self, head, embedding_dim, dropout=0.1):
        super(MultiHeadedAttention, self).__init__()
        assert embedding_dim % head == 0
        # 每个头分配等量的词特性，也就是embed_dim / head个
        self.d_k = embedding_dim // head
        # 传入的头个数
        self.head = head
        self.embedding_dim = embedding_dim

        # 获得线性层对象，通过nn的Linear实例化，它的内部变换矩阵是embedding_dim * embedd
        # q,k,v和最后拼接的矩阵，一共4个
        self.linears = clones(nn.Linear(embedding_dim, embedding_dim), 4)

        # 初始化注意力张量
        self.attn = None

        # 初始化dropout对象
        self.dropout = nn.Dropout(p=dropout)


    def forward(self, query, key, value, mask=None):
        """
        前向反馈函数，4个参数
        :param query:
        :param key:
        :param value:
        :param mask:
        :return:
        """
        if mask is not None:
            mask = mask.unsqueeze(0)
        batch_size = query.size(0)

        # 接下来进入多头处理环节
        query, key, value = [model(x).view(batch_size, -1, self.head, self.d_k).transpose(1, 2)
                             for model, x in zip(self.linears, (query, key, value))]

        # 将每个头的输出传入到注意力层
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)

        # 得到每个头的计算结果是4维张量， 需要进行形状的转换
        # 前面已经将1,2两个维度进行过转置, 在这里要重新转置回来
        # 注意: 经历了transpose()方法后, 必须要使用contiguous方法, 不然无法使用view()方法
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.head * self.d_k)

        # 最后将x输入线性层列表中的最后一个线性层中进行处理, 得到最终的多头注意力结构输出
        return self.linears[-1](x)


class PositionWiseFeedForward(nn.Module):
    """
    前馈前连接层
    """
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionWiseFeedForward, self).__init__()
        self.d_ff = d_ff
        self.dropout = nn.Dropout(dropout)
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        # x: 代表来自上一层的输出
        # 首先将x送入第一个线性层网络, 然后经历relu函数的激活, 再经历dropout层的处理
        # 最后送入第二个线性层
        return self.w2(self.dropout(F.relu(self.w1(x))))


class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()

        # 初始化两个参数张量a2, b2，用于对结果做规范化操作计算
        # 将其用nn.Parameter进行封装, 代表他们也是模型中的参数
        self.a2 = nn.Parameter(torch.ones(features))
        self.b2 = nn.Parameter(torch.zeros(features))

        self.eps = eps

    def forward(self, x):
        # x: 代表上一层网络的输出
        # 首先对x进行最后一个维度上的求均值操作, 同时操持输出维度和输入维度一致
        mean = x.mean(-1, keepdim=True)
        # 接着对x进行字后一个维度上的求标准差的操作, 同时保持输出维度和输入维度一致
        std = x.std(-1, keepdim=True)
        # 按照规范化公式进行计算并返回
        return self.a2 * (x - mean) / (std + self.eps) + self.b2


class SublayerConnection(nn.Module):
    def __init__(self, size, dropout=0.1):
        # size: 代表词嵌入的维度
        # dropout: 进行Dropout操作的置零比率
        super(SublayerConnection, self).__init__()
        # 实例化一个规范化层的对象
        self.norm = LayerNorm(size)
        # 实例化一个dropout对象
        self.dropout = nn.Dropout(p=dropout)
        self.size = size

    def forward(self, x, sublayer):
        # x: 代表上一层传入的张量
        # sublayer: 该子层连接中子层函数
        # 首先将x进行规范化, 然后送入子层函数中处理, 处理结果进入dropout层, 最后进行残差连接
        return x + self.dropout(sublayer(self.norm(x)))


class EncoderLayer(nn.Module):
    def __init__(self, size, self_attn, feed_forward, dropout):
        """

        :param size:词嵌入的维度
        :param self_attn:传入的多头自注意力子层的实例化对象
        :param feed_forward:前馈全连接层实例化对象
        :param dropout:进行dropout操作时的置零比率
        """
        super(EncoderLayer, self).__init__()

        # 将两个实例化对象和参数传入类中
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.size = size

        # 编码器层中有2个子层连接结构, 使用clones函数进行操作
        self.sublayer = clones(SublayerConnection(size, dropout), 2)

    def forward(self, x, mask):
        # x: 代表上一层的传入张量
        # mask: 代表掩码张量
        # 首先让x经过第一个子层连接结构,内部包含多头自注意力机制子层
        # 再让张量经过第二个子层连接结构, 其中包含前馈全连接网络
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)


class Encoder(nn.Module):
    def __init__(self, layer, N):
        # layer: 代表编码器层
        # N: 代表编码器中有几个layer
        super(Encoder, self).__init__()
        # 首先使用clones函数克隆N个编码器层放置在self.layers中
        self.layers = clones(layer, N)
        # 初始化一个规范化层, 作用在编码器的最后面
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        # x: 代表上一层输出的张量
        # mask: 代表掩码张量
        # 让x依次经历N个编码器层的处理, 最后再经过规范化层就可以输出了
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)


class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        """

        :param size: 词嵌入的维度
        :param self_attn: 多头自注意力机制对象
        :param src_attn: 常规注意力机制对象
        :param feed_forward: 前馈全连接层对象
        :param dropout: 置0比率
        """
        super(DecoderLayer, self).__init__()

        # 将参数传入类中
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.dropout = dropout

        # 按照解码器层的结构图, 使用clones函数克隆3个子层连接对象
        self.sublayer = clones(SublayerConnection(size, dropout), 3)

    def forward(self, x, memory, source_mask, target_mask):
        """

        :param x: 上一层输出张量
        :param memory: 编码器语义存储张量
        :param source_mask: 源数据掩码张量
        :param target_mask: 目标数据掩码张量
        :return:
        """
        m = memory

        # 第一步让x经历第一个子层, 多头自注意力机制的子层
        # 采用target_mask, 为了将解码时未来的信息进行遮掩, 比如模型解码第二个字符, 只能看见第一个字符信息
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, target_mask))

        # 第二步让x经历第二个子层, 常规的注意力机制的子层, Q!=K=V
        # 采用source_mask, 为了遮掩掉对结果信息无用的数据
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, source_mask))

        # 第三步让x经历第三个子层, 前馈全连接层
        return self.sublayer[2](x, self.feed_forward)


class Decoder(nn.Module):
    def __init__(self, layer, N):
        """

        :param layer:解码器层对象
        :param N: 复制多少个layer
        """
        super(Decoder, self).__init__()
        # 利用clones函数克隆N个layer
        self.layers = clones(layer, N)
        # 实例化一个规范化层
        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, source_mask, target_mask):
        """

        :param x:输入数据
        :param memory: 输出张量
        :param source_mask: 源数据掩码张量
        :param target_mask: 目标数据掩码张量
        :return:
        """
        for layer in self.layers:
            x = layer(x, memory, source_mask, target_mask)
        return self.norm(x)


class Generator(nn.Module):
    def __init__(self, d_model, vocab_size):
        """

        :param d_model: 词嵌入维度
        :param vocab_size: 代表词表大小
        """
        super(Generator, self).__init__()
        #
        self.project = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        """
        首先得到上一步得到的self.project对x进行线性变化
        然后使用F中已经实现的log_softmax进行softmax处理
        使用log_softmax是因为和我们这个pytorch版本的损失函数实现有关，在其他版本中将被修复
        log_softmax就是对softmax的结果又取了对数，因为对数函数是单调递增的
        因此对最终结果没有影响
        :param x:
        :return:
        """
        return F.log_softmax(self.project(x), dim=-1)


class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder, source_embed, target_embed, generator):
        """

        :param encoder:编码器对象
        :param decoder:解码器对象
        :param source_embed:源数据的嵌入函数
        :param target_embed:目标数据的嵌入函数
        :param generator:输出部分类别生成器对象
        """
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = source_embed
        self.tgt_embed = target_embed
        self.generator = generator

    def forward(self, source, target, source_mask, target_mask):
        """
        在函数中, 将source, source_mask传入编码函数, 得到结果后,
        与source_mask，target，和target_mask一同传给解码函数.
        :param source: 源数据
        :param target: 目标数据
        :param source_mask: 掩码张量
        :param target_mask:
        :return:
        """
        return self.decode(self.encode(source, source_mask), source_mask, target, target_mask)

    def encode(self, source, source_mask):
        return self.encoder(self.src_embed(source), source_mask)

    def decode(self, memory, source_mask, target, target_mask):
        return self.decoder(self.tgt_embed(target), memory, source_mask, target_mask)


def make_model(source_vocab, target_vocab, N=6, d_model=512, d_ff=2048, head=8, dropout=0.1):
    """
    
    :param source_vocab: 源数据的词汇总数
    :param target_vocab: 目标数据的词汇总数
    :param N: 编码器和解码器堆叠的层数
    :param d_model: 词嵌入的维度
    :param d_ff: 前馈全连接层中变换矩阵的维度
    :param head: 多头注意力机制中的头数
    :param dropout: 置零的比率
    :return: 
    """
    c = copy.deepcopy

    # 实例化一个多头注意力的类
    attn = MultiHeadedAttention(head, d_model)
    # 实例化一个前馈全连接层的网络对象
    ff = PositionWiseFeedForward(d_model, d_ff, dropout)
    # 实例化一个位置编码器
    position = PositionalEncoding(d_model, dropout)

    # 实例化模型model,利用的是EncoderDecoder类
    # 编码器的结构里面有2个子层, attention层和前馈全连接层
    # 解码器的结构中有3个子层, 两个attention层和前馈全连接层
    model = EncoderDecoder(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        nn.Sequential(Embeddings(source_vocab, d_model), c(position)),
        nn.Sequential(Embeddings(target_vocab, d_model), c(position)),
        Generator(d_model, target_vocab)
    )
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model


def test01():
    # mask_ = subsequent_mask(15)
    # print(mask_)
    # plt.figure(figsize=(5, 5))
    # plt.imshow(mask_.squeeze(0))
    # plt.show()
    # exit(0)

    d_model = 512
    vocab = 1000
    max_len = 50
    # input_ batch_size=2, seq_len=4
    x = torch.LongTensor([[100, 2, 421, 528], [491, 998, 1, 221]])
    emb = Embeddings(vocab, d_model)
    embr = emb(x)
    print(embr)
    print(embr.shape)
    exit(0)
    pe = PositionalEncoding(d_model, 0.1, max_len)
    pe_result = pe(embr)
    # print('pe_result: ', pe_result)
    # print('pe_result.shape: ', pe_result.shape)

    # draw_plt() # 画图
    # query = key = value = pe_result
    # attn, p_attn = attention(query, key, value)
    # print('attn: ', attn)
    # print('attn shape: ', attn.shape)
    # print('p_attn: ', p_attn)
    # print('p_attn(attention weights) shape: ', p_attn.shape)

    # 多头注意力验证
    # print('-' * 20, '多头注意力验证')
    head = 8
    embed_dim = 512
    dropout = 0.2
    mask = torch.zeros(head, 4, 4)
    # mha = MultiHeadAttention(head, embed_dim, dropout)
    # mha_result = mha(query, key, value, mask)
    # print(mha_result)
    # print('-' * 40)

    # ff = PositionWiseFeedForward(d_model, embed_dim, dropout)
    # ff_result = ff(pe_result)
    # print('ff_result: ', ff_result)

    # eps = 1e-6
    # ln = LayerNorm(embed_dim, eps)
    # ln_result = ln(ff_result)
    # print('ln_result: ', ln_result)

    # sc = SublayerConnection(embed_dim, dropout)
    # self_attn = MultiHeadAttention(head, d_model)
    # sublayer = lambda x: self_attn(x, x, x, mask)
    # sc_result = sc(pe_result, sublayer)
    # print('sc_result: ', sc_result)

    size = 512
    d_ff = 64
    c = copy.deepcopy
    self_attn = MultiHeadedAttention(head, d_model)
    ff = PositionWiseFeedForward(d_model, d_ff, dropout)
    encoder_layer = EncoderLayer(size, c(self_attn), c(ff), dropout)
    # el_result = el(x, mask)
    # print(el_result)
    # print(el_result.shape)

    N = 8
    en = Encoder(encoder_layer, N)
    # en_result = en(x, mask)
    # print(en_result)
    # print(en_result.shape)

    decoder_layer = DecoderLayer(d_model, c(self_attn), c(self_attn), c(ff), dropout)
    source_mask = target_mask = mask
    # memory = en_result
    de = Decoder(decoder_layer, N)
    # de_result = de(x, memory, source_mask, target_mask)
    # print(de_result)
    # print(de_result.shape)

    gen = Generator(d_model, vocab)
    # gen_result = gen(de_result)
    # print(gen_result)
    # print(gen_result.shape)

    source_embed = nn.Embedding(vocab, d_model)
    target_embed = nn.Embedding(vocab, d_model)
    ende = EncoderDecoder(en, de, source_embed, target_embed, gen)
    ed_result = ende(x, x, source_mask, target_mask)
    print(ed_result)
    print(ed_result.shape)

    source_vocab = 11
    target_vocab = 11
    N = 6
    res = make_model(source_vocab, target_vocab, N)
    print(res)


if __name__ == '__main__':
    test01()
    pass